
#include <openssl/evp.h>
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "../iface/SSLInterface.h"
#include "AESEncrypt.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe encryption.
 * 
 * @param aeskey The AES key to use for encryption. */
AESEncrypt::AESEncrypt(const AESSymmetricKey &aeskey)
{
   const char *aes_keypass = aeskey.getKeyData();
   keypassword = (aes_keypass) ? strdup(aes_keypass) : NULL;
   input = NULL;
   output = NULL;
   if (keypassword == NULL)
   {
      errprint("AES key is NULL");
      binarykey = NULL;
      binarylength = 0;
   } else
   {
      binarykey = BaseConvert::convertBase64ToBytes(keypassword, &binarylength);
   }
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe encryption.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output.
 * 
 * @param aeskey The AES key to use for encryption. */
AESEncrypt::AESEncrypt(APRBucketBrigade &iinput, APRBucketBrigade &ioutput, const AESSymmetricKey &aeskey)
{
   const char *aes_keypass = aeskey.getKeyData();
   keypassword = (aes_keypass) ? strdup(aes_keypass) : NULL;
   input = &iinput;
   output = &ioutput;
   if (keypassword == NULL)
   {
      errprint("AES key is NULL");
      binarykey = NULL;
      binarylength = 0;
   } else
   {
      binarykey = BaseConvert::convertBase64ToBytes(keypassword, &binarylength);
   }
}

/** Destructor for the AES encryption class. */
AESEncrypt::~AESEncrypt()
{
   StringUtils::razeString(&keypassword);
   memset(binarykey, 0, sizeof(char)*binarylength);
   binarylength = 0;
   if (binarykey) free(binarykey);
}

/** Encrypt data block and return encrypted data. The encrypted data
 * returned must be freed by the caller.
 * 
 * @param plainblock A block of plain data to encrypt.
 * 
 * @param plainsize The length of the plain data block.
 * 
 * @param[out] codedsize The length of the coded block of data
 * returned. The data returned must be freed by the caller. */
char *AESEncrypt::processDataBlock(const char *plainblock, apr_size_t plainsize, apr_size_t *codedsize)
{
   EVP_CIPHER_CTX ctxt;
   EVP_CIPHER_CTX_init(&ctxt);
   int bytes = strlen(keypassword);
   unsigned char *binarykey = (unsigned char *)BaseConvert::convertBase64ToBytes(keypassword, &bytes);
   EVP_CipherInit_ex(&ctxt, EVP_aes_256_cbc(), NULL, NULL, NULL, 1);
   EVP_CIPHER_CTX_set_key_length(&ctxt, (AES_BITS/8));
   unsigned char *startiv = (unsigned char *)calloc(bytes, sizeof(unsigned char));
   if (bytes > (AES_BITS / 8))
      memcpy(startiv, &(binarykey[AES_BITS / 8]), sizeof(char)*(bytes - (AES_BITS/8)));
   char *result = (char *)malloc(sizeof(char)*(plainsize + EVP_CIPHER_CTX_block_size(&ctxt)));
   EVP_CipherInit_ex(&ctxt, EVP_aes_256_cbc(), NULL, binarykey, startiv, 1);
   int outlen = -1;
   int outidx = 0;
   if (!EVP_CipherUpdate(&ctxt, (unsigned char *)(&(result[outidx])),
      &outlen, (const unsigned char *)plainblock, plainsize))
   {
      errprint("failed updating cipher");
      SSLCheckError();
      goto cip_cleanup;
   }
   outidx += outlen;
   outlen = -1;
   if (!EVP_CipherFinal_ex(&ctxt, (unsigned char *)(&(result[outidx])), &outlen))
   {
      errprint("failed finalising cipher");
      SSLCheckError();
      goto cip_cleanup;
   }
   if (outidx > (plainsize + EVP_CIPHER_CTX_block_size(&ctxt)))
   {
      errprint("overflowed cipher (expecting at most %" APR_SIZE_T_FMT ", got %" APR_SIZE_T_FMT ")", (plainsize + EVP_CIPHER_CTX_block_size(&ctxt)), (apr_size_t)outidx);
      goto cip_cleanup;
   }
   if (codedsize)
      *codedsize = outidx + outlen;
   EVP_CIPHER_CTX_cleanup(&ctxt);
   free(startiv); free(binarykey);
   return result;
cip_cleanup:
   EVP_CIPHER_CTX_cleanup(&ctxt);
   free(startiv); free(binarykey); free(result);
   return NULL;
}

/** Begin the process of encrypting a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int AESEncrypt::processPipeBegin()
{
   EVP_CIPHER_CTX_init(&e_ctxt);
   int bytes = strlen(keypassword);
   unsigned char *binarykey = (unsigned char *)BaseConvert::convertBase64ToBytes(keypassword, &bytes);
   if (!EVP_CipherInit_ex(&e_ctxt, EVP_aes_256_cbc(), NULL, NULL, NULL, 1))
   {
      errprint("could not initialise cipher");
      free(binarykey);
      return 1;
   }
   EVP_CIPHER_CTX_set_key_length(&e_ctxt, (AES_BITS/8));
   unsigned char *startiv = (unsigned char *)calloc(bytes, sizeof(unsigned char));
   if (bytes > (AES_BITS / 8))
      memcpy(startiv, &(binarykey[AES_BITS / 8]), sizeof(char)*(bytes - (AES_BITS/8)));
   if (!EVP_CipherInit_ex(&e_ctxt, EVP_aes_256_cbc(), NULL, binarykey, startiv, 1))
   {
      errprint("could not initialise cipher with key");
      free(startiv);
      free(binarykey);
      return 1;
   }
   free(startiv);
   free(binarykey);
   return 0;
}

/** Process any outstanding plaintext data on the input pipe, and
 * place any processed encrypted data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int AESEncrypt::processPipeStep()
{
   int cip_blocksize = EVP_CIPHER_CTX_block_size(&e_ctxt);
   int input_block = input->realLength();
   char *itemp_block = (char *)malloc(sizeof(char)*input_block);
   char *otemp_block = (char *)malloc(sizeof(char)*(input_block + cip_blocksize));
   
   apr_size_t readbytes = input->readFromStream(itemp_block, input_block);
   // this isn't a problem, this should fail during normal operation
   // when the buffer runs out of data.
   if (readbytes != input_block)
   {
      errprint("failed to read from stream");
      free(itemp_block);
      free(otemp_block);
      return 1;
   }        
   int outlen = -1;
   if (!EVP_CipherUpdate(&e_ctxt, (unsigned char *)otemp_block,
      &outlen, (const unsigned char *)itemp_block, input_block))
   {
      errprint("failed updating cipher");
      SSLCheckError();
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   apr_size_t writebytes = output->writeToStream(otemp_block, outlen);
   if (writebytes != outlen)
   {
      errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   input->removeFromStream(input_block);
   free(itemp_block);
   free(otemp_block);
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int AESEncrypt::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding plaintext data in the pipe, process any
 * outstanding data in the internal buffers of the encryption process
 * and finalise the encryption process.
 * 
 * @return Zero on success, non-zero on failure (one if the cipher has
 * not been finalised and this can be run again, two of the cipher has
 * been finalised but failed). */
int AESEncrypt::processPipeEnd()
{
   processPipeStep();
   int cip_blocksize = EVP_CIPHER_CTX_block_size(&e_ctxt);
   apr_size_t plaindata_blocksize = input->realLength();
   if (plaindata_blocksize >= cip_blocksize)
   {
      errprint("too much data remaining in bucket brigade (is it still being written to?)");
      return 1;
   }
   char *itemp_block = (char *)malloc(sizeof(char)*plaindata_blocksize);
   char *otemp_block = (char *)malloc(sizeof(char)*(plaindata_blocksize + cip_blocksize));

   apr_size_t readbytes = input->readFromStream(itemp_block, plaindata_blocksize);
   if (readbytes != plaindata_blocksize)
   {
      errprint("failed to read data from plain bucket brigade (read %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), buffer unfinished", readbytes, plaindata_blocksize);
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   int outlen = -1;
   if (!EVP_CipherUpdate(&e_ctxt, (unsigned char *)otemp_block,
      &outlen, (const unsigned char *)itemp_block, plaindata_blocksize))
   {
      errprint("failed updating cipher");
      SSLCheckError();
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   apr_size_t writebytes = output->writeToStream(otemp_block, outlen);
   if (writebytes != outlen)
   {
      errprint("failed to write data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
      free(itemp_block);
      free(otemp_block);
      return 1;
   }
   input->removeFromStream(plaindata_blocksize);
   outlen = -1;
   if (!EVP_CipherFinal_ex(&e_ctxt, (unsigned char *)(otemp_block), &outlen))
   {
      errprint("failed finalising cipher");
      free(itemp_block);
      free(otemp_block);
      SSLCheckError();
      EVP_CIPHER_CTX_cleanup(&e_ctxt);
      return 2;
   }
   writebytes = output->writeToStream(otemp_block, outlen);
   if (writebytes != outlen)
   {
      errprint("failed to write final data to coded bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", writebytes, (apr_size_t)(outlen));
      free(itemp_block);
      free(otemp_block);
      EVP_CIPHER_CTX_cleanup(&e_ctxt);
      return 2;
   }
   free(itemp_block);
   free(otemp_block);
   EVP_CIPHER_CTX_cleanup(&e_ctxt);
   return 0;
}
